[pull] master from deepmodeling:master - #316

…through pt_expt assembly (#5960)

Users reported that compiled DPA4 training runs ~2x slower on `pt_expt`
than on `pt`. This PR is the result of chasing that: a performance bug
in how the SO3 contractions were lowered, and a correctness bug where a
configured `use_amp: false` was lost while `pt_expt` assembled the
model.

## Changes

**1. Weight broadcast across the node axis (`so3.py`, `lora.py`,
`grid_net.py`)**

`matmul(x[..., None, :], weight[None, ...])` makes the node count `N`
the matmul BATCH, so matmul broadcasts the weight to `(N, D, F, Cin,
Cout)` and autograd then reduces that whole expanded gradient
(`ExpandBackward0`) back to the parameter shape. At the water example's
sizes a 165 K-element weight expanded to 191 M elements (~0.8 GB) per
call, and the reduce was the single costliest kernel of a training step
(45.6 ms, 3x per step).

The fix batches the contraction over the small `(D, F)` axes so `N`
stays the GEMM ROW dimension and the weight is used in place.
Micro-benchmark, fwd+bwd at the real shapes: **16.48 ms -> 1.09 ms
(15x)**.

Two lookalike sites in `projection.py` are deliberately NOT changed:
their operands are `requires_grad=False` buffers, so no backward reduce
exists. Verified rather than assumed.

**2. The same lowering for the frame mixers (`_degree_batched_matmul`)**

Review found the `FrameContract` / `FrameExpand` mixers still on the
broadcast spelling. They now share one helper, `_degree_batched_matmul`,
written identically on the dpmodel side
(`dpmodel/descriptor/dpa4_nn/grid_net.py`) and the pt side
(`pt/model/descriptor/sezm_nn/grid_net.py`). Because it does no reshape,
an empty node axis (`N == 0`) flows through unchanged instead of hitting
a reshape error — pinned by a test.

**3. `use_amp` was lost during `pt_expt` model assembly (correctness)**

`use_amp` is a training-runtime policy, not model state, so it stays OUT
of the portable serialization record (this is the #5963 position, and
the jax deserializer actively rejects records carrying `use_amp: true`).
The bug was elsewhere: `pt_expt` assembled its model by converting an
already-populated dpmodel instance, and that conversion round-trips the
component through `deserialize(serialize())`. Anything that is
deliberately not in the portable record — `use_amp` among it — was
therefore dropped, and training silently ran under bfloat16 autocast
even when the input configured `use_amp: false`.

The fix is at the assembly boundary, not in the record: `pt_expt` now
constructs the wrapped class directly (`auto_wrapped_class(...)` in
`make_model.py`, `get_model.py`, and the bridging composition path), so
a live constructor-supplied component keeps its runtime state. The rule
is stated once, in the `auto_wrapped_class` docstring; the call sites
reference it.

An earlier revision of this PR instead added `use_amp` to `serialize()`.
That was reverted in review — it put a runtime knob into the portable
record and would have broken the jax contract.

**Also removed in review: an `enable_tf32` / `DP_TF32_INFER`
implementation for `pt_expt`.** It contributes nothing to the speedup
measured below (the benchmark card has no TF32 silicon), and #5958 owns
the `pt_expt` training-runtime alignment — including the documented
position that `pt_expt` always runs at `"highest"` matmul precision.
`pt_expt` therefore keeps master's warn-and-ignore behavior for
`enable_tf32`.

## Benchmark

DPA4 water example (`examples/water/dpa4`), one Tesla T4, torch 2.11,
fp32 (`use_amp: false`), batch size 6. Steady-state seconds per training
step, obtained by differencing the wall time of a 33-step and a 3-step
run of the same config, which cancels every one-time cost (import, data
load, statistics, `torch.compile` / make_fx lowering). All five arms
were measured in one session on the same machine; run-to-run variation
is about 2-3%.

**Provenance: measured at `ae720432b`**, the head at which this PR was
opened — i.e. BEFORE the review changes (change 2, the frame-mixer
lowering, and change 3's move from `serialize()` to the assembly
boundary). Change 1, which is where the entire speedup comes from, is
unmodified since. The numbers have not been re-measured on the current
head; a re-run is pending and I will post it rather than silently reuse
these.

| training mode | `pt` (reference) | `pt_expt` at master | `pt_expt` at
`ae720432b` | speedup vs master |
|---|---|---|---|---|
| eager | 0.891 s/step | 1.555 s/step | **0.921 s/step** | **1.69x** |
| compiled | 0.545 s/step | 1.611 s/step | **0.535 s/step** | **3.01x**
|

This reproduces the reported issue at master — `pt_expt` compiled was
3.0x slower than `pt` compiled, and even slower than its own eager path,
because the broadcast-weight contraction lowers to worse code under
inductor than under eager cuBLAS. After the fix `pt_expt` is at parity
with `pt`: eager within 3.4%, compiled within measurement noise.

## Known limitations

- **Backward numerics are covered for the frame mixers, not for the SO3
/ LoRA contractions.** `test_dpa4_frame_mixers.py` compares
`_degree_batched_matmul`'s weight gradient against the pt module's at
rtol/atol 1e-12. The rewritten SO3 and LoRA contractions are pinned on
the forward against an explicit `einsum` reference (rtol/atol 1e-12,
numpy and torch namespaces); their backward is still exercised only by
tracing, not compared by value.
- **The `pt` / `pt_expt` TF32 policy gap remains open.** On Ampere+
cards `pt` runs training matmuls under TF32 (`enable_tf32`, default
`True`) while `pt_expt` ignores the key with a warning; the two backends
are not speed-comparable there. Deferred to the #5958 training-runtime
series.
- **The residual compiled gap vs `pt` is not stable across sessions.**
An earlier session measured `pt_expt` compiled 10.8% slower than `pt`
compiled; the benchmark above measured it 1.7% faster. Both are within a
couple of run-to-run standard deviations, so I treat compiled as at
parity and the earlier gap as unconfirmed.
- The history contains churn at the GridBranch router (`7518a417c` ->
`01c58e665` -> `75459610a` -> `504bb2430` -> `157444204`): a matmul
spelling introduced, reverted, reintroduced, and finally restored to
master's line. The site is byte-identical to master in the final diff.
The degenerate GEMM that profiling found there existed only on this
branch, so it is not a fix — I have left the commits rather than
rewriting pushed history, and would squash them on request.
- Unrelated but found while benchmarking: **torch >= 2.11 ships no Volta
(CC 7.0) kernels**, and compiled training requires >= 2.11 via
`check_compile_torch_version`. Compiled DPA4 training is therefore
impossible on V100 with official wheels; T4 (CC 7.5) is the oldest card
that works.

## Tests

- `source/tests/common/dpmodel/test_dpa4_frame_mixers.py` —
`_degree_batched_matmul` vs the pt module: forward parity, the `N == 0`
contract, and weight-gradient parity.
- `source/tests/common/dpmodel/test_dpa4_lora.py` — new
`test_lora_so3_call_matches_einsum_contract`: `LoRASO3.call` against the
explicit `einsum("ndfi,difo->ndfo")` reference with a nonzero adapter,
on both the numpy and torch namespaces, for `n_focus` 1 and 2.
- `source/tests/pt_expt/model/test_get_model_dpa4.py` — `use_amp`
survives model assembly (both branches), for the plain and the
bridged/composed construction paths.
- `source/tests/common/dpmodel/test_descrpt_dpa4.py` — `use_amp` is
absent from the portable serialization record and defaults on
deserialize.
- Existing `test_grid_branch[1]`/`[2]` cover the changed SO3 contraction
against the pt implementation at rtol 1e-12.
- Run locally: 434 passed / 10 skipped across the dpa4 dpmodel, pt_expt
and cross-backend parity suites, plus the pt_expt model suite.
CUDA-gated precision-context cases were run on a T4 (29/29).

## Test status caveat — resolved

An earlier revision of this description flagged two locally failing
`pt_expt` AOTI-freeze tests
(`test_zbl_bridging.py::test_native_spin_with_bridging_graph_freeze_and_deep_eval`,
`test_dpa4_zbl_parallel.py::TestBridgedSpinGraphSelfComm::test_freeze_embeds_with_comm_artifact`)
as unadjudicated. They are now adjudicated as **pre-existing and
environmental, not caused by this branch**: a clean `upstream/master`
worktree on the same machine fails both with the identical
`InductorError: assert isinstance(index, CppCSEVariable) and
index.is_vec` (torch 2.11 CPU-SIMD codegen bug on an `atomic_add`
scatter buffer), and both tests pass on this branch with the known
workaround `torch._inductor.config.cpp.simdlen = 1` (2 passed). The same
bug is already documented in `source/tests/infer/gen_dpa4.py` /
`gen_dpa2.py`.

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>